package com.feidee.fd.sml.algorithm.component.ml

import com.feidee.fd.sml.algorithm.component.{AbstractComponent, BasicParam}
import org.apache.spark.ml.{PipelineModel, PipelineStage}
import org.apache.spark.sql.{DataFrame, functions}

/**
  * @Author songhaicheng
  * @Date 2018/08/27
  * @Email: haicheng_song@sui.com
  */
trait MLParam extends BasicParam {

  override val input_pt: String
  override val output_pt: String
  override val hive_table: String
  override val flow_time: String

  // 特征列名，默认 features
  val featuresCol: String
  // 标签列名，默认 label
  var labelCol: String
  // 预测值列名，默认 prediction
  var predictionCol: String
  // 模型保存路径，为空时不保存
  val modelPath: String
  // 需要获取的评估指标（分类: accuracy, f1, auc, pr，回归: mse, rmse, r2, mae）
  val metrics: Array[String]

  override def verify(): Unit = {
    super.verify()
    require(tool.isNotNull(input_pt), "param input_pt can't be null")
    require(tool.isNotNull(featuresCol), "输入特征字段 featuresCol 不能为空")
    require(tool.isNotNull(labelCol), "标签字段 labelCol 不能为空")
    require(tool.isNotNull(predictionCol), "预测结果字段 predictionCol 不能为空")
    require(!featuresCol.equals(labelCol) && !featuresCol.equals(predictionCol),
      "特征字段 featuresCol、标签字段 labelCol、结果字段 predictionCol 不能出现一样")
    require(tool.isNotNull(modelPath),"modelPath 不能为空")
  }

  override def toMap: Map[String, Any] = {
    var map = super.toMap
    map += ("featuresCol" -> featuresCol)
    map += ("labelCol" -> labelCol)
    map += ("predictionCol" -> predictionCol)
    map += ("modelPath" -> modelPath)
    map += ("metrics" -> metrics)
    map
  }

}


abstract class AbstractMLComponent[A <: MLParam] (implicit m:Manifest[A]) extends AbstractComponent[A] with Serializable {

  /**
    * 根据参入参数设置算法组件实体
    * @param param 运行参数
    * @return      算法组件
    */
  def setUp(param: A): PipelineStage = ???

  /**
    * 训练模型
    * @param param 算法运行参数
    * @param data  输入（训练用）数据
    * @return      算法模型
    */
  def train(param: A, data: DataFrame): PipelineModel = ???

  /**
    * 根据参数，评估指标
    * @param param    模型训练时的参数
    * @param model    生成的模型（可以从模型里拿相关信息辅助指标计算）
    * @param data     评估用的 DataFrame
    * @return         指标值
    */
  def calculateMetrics(param: A, model: PipelineModel, data: DataFrame): Map[String, Double] = ???

  /**
    * 保存模型，除了要保存一份到参数指定路径，还需要在该路径下保存一份以时间作为后缀的模型作为备份及供线上使用
    * @param model  要保存的模型
    * @param param  参数
    */
  def outputModel(model: PipelineModel, param: A): Unit = {
    model.write.overwrite().save(param.modelPath)
    if (tool.isNotNull(param.flow_time)) {
      model.write.overwrite().save(s"${param.modelPath}_${param.flow_time}")
    } else {
      logWarning("未发现运行时间参数，不做模型备份处理")
    }
  }

  /**
    * 保存指标
    * @param metrics  保存的指标集合
    * @param dir      指标文件存放目录
    * @return         是否都保存成功
    */
  @Deprecated
  def outputMetrics(metrics: Map[String, Double], dir: String): Boolean = {
    var isMetricsAllSaved = true
    // 在存放目录下写指标值入对应指标文件名
    metrics.foreach{case (name, value) =>
      isMetricsAllSaved &= tool.writeToHDFS(s"$dir/$name", value.toString)
    }
    isMetricsAllSaved
  }

  /**
    * 把计算的指标添加到结果 DataFrame 后（表地址就是设置的 hive_table）
    * @param metrics  指标值
    * @param data     结果 DataFrame
    * @return         追加指标值后的结果 DataFrame
    */
  def appendMetrics(metrics: Map[String, Double], data: DataFrame): DataFrame = {
    metrics.foldLeft(data)((cur, m) => cur.withColumn(m._1, functions.lit(m._2)))
  }

  /**
    * 执行过程
    * @param paramStr 传过来的参数
    */
  override def apply(paramStr: String): Unit = {
    logInfo("parsing parameter")
    val param = parseParam(paramStr)
    logInfo("verifying parameter")
    param.verify()
    logInfo(s"loading input data FROM ${param.input_pt}")
    val inputData = loadData(param)
    logInfo("training ml model")
    val model = train(param, inputData)
    var result = model.transform(inputData)
    // 保存
    if (tool.isNotNull(param.output_pt)) {
      logInfo(s"saving ml predicted result TO ${param.output_pt}")
      outputData(result, param)
    }
    if (tool.isNotNull(param.modelPath)) {
      logInfo(s"saving ml model TO ${param.modelPath}")
      outputModel(model, param)
    }
    // 另存
    if (tool.isNotNull(param.hive_table)) {
      // 要计算的指标不为空
      if (param.metrics.length > 0) {
        logInfo("calculating metrics")
        val metrics = calculateMetrics(param, model, result)
        // 追加指标值到结果中
        logInfo(s"appending metrics ${metrics.keySet.mkString("[", ", ", "]")}")
        result = appendMetrics(metrics, result)
      }
      logInfo(s"saving ml predicted result to hive table ${param.hive_table}")
      outputTable(result, param)
    }
  }

}
